package SingleCycleProcessor

import chisel3._
import chisel3.util._
import scala.io.Source
import common.Instructions._
import common.Consts._



//------------Memory-------------

class ImemPortIo extends Bundle {
  val addr = Input(UInt(WORD_LEN.W))
  val inst = Output(UInt(WORD_LEN.W))
}

class DmemPortIo extends Bundle {
  val addr  = Input(UInt(WORD_LEN.W))
  val rdata = Output(UInt(WORD_LEN.W))
  val wen   = Input(Bool())
  val wdata = Input(UInt(WORD_LEN.W))
}

class Memory extends Module {
  val io = IO(new Bundle {
    val imem = new ImemPortIo()
    val dmem = new DmemPortIo()
  })

  val mem = Mem(16384, UInt(8.W))
  //loadMemoryFromFile(mem, "src/riscv/rv32ui-p-sw.hex")
  var fileLines = Source.fromFile("src/hex/rv32ui-p-sw.hex").getLines().toList
  for((line, index) <- fileLines.zipWithIndex) {
    val num: Int = Integer.parseInt(line, 16)
    val idxBits: UInt = index.asUInt
    val lineBits: UInt = num.asUInt
    mem(index) := lineBits
    //chisel3.printf(p"fromFile: Index: ${idxBits}, Content: ${lineBits}, mem: ${mem(index)}\n")
  }
  mem(96.U) := 10.U(WORD_LEN.W)
  mem(100.U) := 11.U(WORD_LEN.W)
  mem(104.U) := 22.U(WORD_LEN.W)
  mem(108.U) := 33.U(WORD_LEN.W)
  mem(112.U) := 44.U(WORD_LEN.W)
  mem(116.U) := 55.U(WORD_LEN.W)

  io.imem.inst := Cat(
    mem(io.imem.addr + 3.U(WORD_LEN.W)), 
    mem(io.imem.addr + 2.U(WORD_LEN.W)),
    mem(io.imem.addr + 1.U(WORD_LEN.W)),
    mem(io.imem.addr)
  )
  io.dmem.rdata := Cat(
    mem(io.dmem.addr + 3.U(WORD_LEN.W)),
    mem(io.dmem.addr + 2.U(WORD_LEN.W)), 
    mem(io.dmem.addr + 1.U(WORD_LEN.W)),
    mem(io.dmem.addr)
  )
  chisel3.printf("io.imem.addr=0x%x,  io.imem.inst=0x%x, io.dmem.addr=0x%x, io.dmem.rdata=0x%x.\n",
                 io.imem.addr,  io.imem.inst, io.dmem.addr, io.dmem.rdata)

  when(io.dmem.wen){
    mem(io.dmem.addr)                   := io.dmem.wdata( 7,  0)
    mem(io.dmem.addr + 1.U(WORD_LEN.W)) := io.dmem.wdata(15,  8)
    mem(io.dmem.addr + 2.U(WORD_LEN.W)) := io.dmem.wdata(23, 16)
    mem(io.dmem.addr + 3.U(WORD_LEN.W)) := io.dmem.wdata(31, 24)
    chisel3.printf("io.dmem.addr=0x%x, io.dmem.wdata=0x%x.\n",io.dmem.addr, io.dmem.wdata)
  }
}

//------------Core---------------
class Core extends Module {
  val io = IO(
    new Bundle {
      val imem = Flipped(new ImemPortIo())
      val dmem = Flipped(new DmemPortIo())
      val exit = Output(Bool())
      val gp   = Output(UInt(WORD_LEN.W))
    }
  )

  val regfile     = Mem(32,   UInt(WORD_LEN.W))
  val csr_regfile = Mem(4096, UInt(WORD_LEN.W)) 
  regfile(3) := 1.U(WORD_LEN.W) 
  //regfile(4) := 0.U(WORD_LEN.W) 
  //regfile(6) := 11.U(WORD_LEN.W) 
  //regfile(8) := 22.U(WORD_LEN.W) 
  //regfile(10) := 33.U(WORD_LEN.W) 
  //regfile(12) := 44.U(WORD_LEN.W) 
  //regfile(14) := 55.U(WORD_LEN.W) 
  regfile(24) := 96.U(WORD_LEN.W) 

  //regfile(1) := 0x01ffff00.U(WORD_LEN.W)


  //**********************************
  // Instruction Fetch (IF) Stage
  
  val pc_reg = RegInit(START_ADDR)
  io.imem.addr := pc_reg
  val inst = io.imem.inst
  val pc_plus4 = pc_reg + 4.U(WORD_LEN.W)
  val br_target = Wire(UInt(WORD_LEN.W))
  val br_flg = Wire(Bool())
  val jmp_flg = (inst === JAL || inst === JALR)
  val alu_out = Wire(UInt(WORD_LEN.W))
  
  val pc_next = MuxCase(pc_plus4, Seq(
    br_flg  -> br_target,
    jmp_flg -> alu_out,
    (inst === ECALL) -> csr_regfile(0x305) // go to trap_vector
  ))
  pc_reg := pc_next


  //**********************************
  // Instruction Decode (ID) Stage

  val rs1_addr = inst(19, 15)
  val rs2_addr = inst(24, 20)
  val wb_addr  = inst(11, 7)
  val rs1_data = Mux((rs1_addr =/= 0.U(WORD_LEN.U)), regfile(rs1_addr), 0.U(WORD_LEN.W))
  val rs2_data = Mux((rs2_addr =/= 0.U(WORD_LEN.U)), regfile(rs2_addr), 0.U(WORD_LEN.W))

  val imm_i = inst(31, 20)
  val imm_i_sext = Cat(Fill(20, imm_i(11)), imm_i) //LW instruction is I-type
  val imm_s = Cat(inst(31, 25), inst(11, 7))
  val imm_s_sext = Cat(Fill(20, imm_s(11)), imm_s)
  val imm_b = Cat(inst(31), inst(7), inst(30, 25), inst(11, 8))
  val imm_b_sext = Cat(Fill(19, imm_b(11)), imm_b, 0.U(1.U))
  val imm_j = Cat(inst(31), inst(19, 12), inst(20), inst(30, 21))
  val imm_j_sext = Cat(Fill(11, imm_j(19)), imm_j, 0.U(1.U))
  val imm_u = inst(31,12)
  val imm_u_shifted = Cat(imm_u, Fill(12, 0.U))
  val imm_z = inst(19,15)
  val imm_z_uext = Cat(Fill(27, 0.U), imm_z)

  val csignals = ListLookup(inst,
               List(ALU_X    , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_00),
    Array(
      LW    -> List(ALU_ADD  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_MEM, CSR_X, INST_01),
      SW    -> List(ALU_ADD  , OP1_RS1, OP2_IMS, MEN_S, REN_X, WB_X  , CSR_X, INST_02),
      ADD   -> List(ALU_ADD  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_03),
      ADDI  -> List(ALU_ADD  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_04),
      SUB   -> List(ALU_SUB  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_05),
      AND   -> List(ALU_AND  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_06),
      OR    -> List(ALU_OR   , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_07),
      XOR   -> List(ALU_XOR  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_08),
      ANDI  -> List(ALU_AND  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_09),
      ORI   -> List(ALU_OR   , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_10),
      XORI  -> List(ALU_XOR  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_11),
      SLL   -> List(ALU_SLL  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_12),
      SRL   -> List(ALU_SRL  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_13),
      SRA   -> List(ALU_SRA  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_14),
      SLLI  -> List(ALU_SLL  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_15),
      SRLI  -> List(ALU_SRL  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_16),
      SRAI  -> List(ALU_SRA  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_17),
      SLT   -> List(ALU_SLT  , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_18),
      SLTU  -> List(ALU_SLTU , OP1_RS1, OP2_RS2, MEN_X, REN_S, WB_ALU, CSR_X, INST_19),
      SLTI  -> List(ALU_SLT  , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_20),
      SLTIU -> List(ALU_SLTU , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_ALU, CSR_X, INST_21),
      BEQ   -> List(BR_BEQ   , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_22),
      BNE   -> List(BR_BNE   , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_23),
      BGE   -> List(BR_BGE   , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_24),
      BGEU  -> List(BR_BGEU  , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_25),
      BLT   -> List(BR_BLT   , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_26),
      BLTU  -> List(BR_BLTU  , OP1_RS1, OP2_RS2, MEN_X, REN_X, WB_X  , CSR_X, INST_27),
      JAL   -> List(ALU_ADD  , OP1_PC , OP2_IMJ, MEN_X, REN_S, WB_PC , CSR_X, INST_28),
      JALR  -> List(ALU_JALR , OP1_RS1, OP2_IMI, MEN_X, REN_S, WB_PC , CSR_X, INST_29),
      LUI   -> List(ALU_ADD  , OP1_X  , OP2_IMU, MEN_X, REN_S, WB_ALU, CSR_X, INST_30),
      AUIPC -> List(ALU_ADD  , OP1_PC , OP2_IMU, MEN_X, REN_S, WB_ALU, CSR_X, INST_31),
      CSRRW -> List(ALU_COPY1, OP1_RS1, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_W, INST_32),
      CSRRWI-> List(ALU_COPY1, OP1_IMZ, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_W, INST_33),
      CSRRS -> List(ALU_COPY1, OP1_RS1, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_S, INST_34),
      CSRRSI-> List(ALU_COPY1, OP1_IMZ, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_S, INST_35),
      CSRRC -> List(ALU_COPY1, OP1_RS1, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_C, INST_36),
      CSRRCI-> List(ALU_COPY1, OP1_IMZ, OP2_X  , MEN_X, REN_S, WB_CSR, CSR_C, INST_37),
      ECALL -> List(ALU_X    , OP1_X  , OP2_X  , MEN_X, REN_X, WB_X  , CSR_E, INST_38)
		)
	)
  val exe_fun :: op1_sel :: op2_sel :: mem_wen :: rf_wen :: wb_sel :: csr_cmd :: inst_code :: Nil = csignals

  val op1_data = MuxCase(0.U(WORD_LEN.W), Seq(
    (op1_sel === OP1_RS1) -> rs1_data, //ADD LW SW are all OP1_RS1 <value = 0>
    (op1_sel === OP1_PC)  -> pc_reg,
    (op1_sel === OP1_IMZ) -> imm_z_uext
  ))
  printf(p"******** op1_sel = ${op1_sel}, op1_data = ${op1_data} !!!\n")

  val op2_data = MuxCase(0.U(WORD_LEN.W), Seq(
    (op2_sel === OP2_RS2) -> rs2_data, //ADD instructon is R-type. <value = 1>
    (op2_sel === OP2_IMI) -> imm_i_sext, //LW instruction is I-type. <value =2>
    (op2_sel === OP2_IMS) -> imm_s_sext, //SW instruction is S-type. <value = 3>
    (op2_sel === OP2_IMJ) -> imm_j_sext,
    (op2_sel === OP2_IMU) -> imm_u_shifted
  ))
  printf(p"******** op2_sel = ${op2_sel}, op2_data = ${op2_data} !!!\n")


  //**********************************
  // Execute (EX) Stage

  alu_out := MuxCase(0.U(WORD_LEN.W), Seq(
    (exe_fun === ALU_ADD)   -> (op1_data + op2_data),
    (exe_fun === ALU_SUB)   -> (op1_data - op2_data),
    (exe_fun === ALU_AND)   -> (op1_data & op2_data),
    (exe_fun === ALU_OR)    -> (op1_data | op2_data),
    (exe_fun === ALU_XOR)   -> (op1_data ^ op2_data),
    (exe_fun === ALU_SLL)   -> (op1_data << op2_data(4, 0))(31, 0),
    (exe_fun === ALU_SRL)   -> (op1_data >> op2_data(4, 0)).asUInt,
    (exe_fun === ALU_SRA)   -> (op1_data.asSInt >> op2_data(4, 0)).asUInt,
    (exe_fun === ALU_SLT)   -> (op1_data.asSInt < op2_data.asSInt).asUInt,
    (exe_fun === ALU_SLTU)  -> (op1_data < op2_data).asUInt,
    (exe_fun === ALU_JALR)  -> ((op1_data + op2_data) & ~1.U(WORD_LEN.W)),
    (exe_fun === ALU_COPY1) -> op1_data
  ))
  printf(p"******** exe_fun = ${exe_fun}, alu_out = ${alu_out} !!!\n")

  // branch
  br_target := pc_reg + imm_b_sext
  br_flg := MuxCase(false.B, Seq(
    (exe_fun === BR_BEQ)  ->  (op1_data === op2_data),
    (exe_fun === BR_BNE)  -> !(op1_data === op2_data),
    (exe_fun === BR_BLT)  ->  (op1_data.asSInt < op2_data.asSInt),
    (exe_fun === BR_BGE)  -> !(op1_data.asSInt < op2_data.asSInt),
    (exe_fun === BR_BLTU) ->  (op1_data < op2_data),
    (exe_fun === BR_BGEU) -> !(op1_data < op2_data)
  ))


  //**********************************
  // Memory Access Stage

  io.dmem.addr  := alu_out
  io.dmem.wen   := mem_wen
  io.dmem.wdata := rs2_data
  when(inst_code === INST_02) {
    printf(p"Store(SW) wen/addr/wdata = ${io.dmem.addr} ${mem_wen} ${rs2_data}!!!\n")
  }

  // CSR
  val csr_addr = Mux(csr_cmd === CSR_E, 0x342.U(CSR_ADDR_LEN.W), inst(31,20))
  val csr_rdata = csr_regfile(csr_addr)
  val csr_wdata = MuxCase(0.U(WORD_LEN.W), Seq(
    (csr_cmd === CSR_W) -> op1_data,
    (csr_cmd === CSR_S) -> (csr_rdata | op1_data),
    (csr_cmd === CSR_C) -> (csr_rdata & ~op1_data),
    (csr_cmd === CSR_E) -> 11.U(WORD_LEN.W)
  ))
  
  when(csr_cmd > 0.U){
    csr_regfile(csr_addr) := csr_wdata
  }

  //**********************************
  // Writeback (WB) Stage
  
  val wb_data = MuxCase(alu_out, Seq(
    (wb_sel === WB_MEM) -> io.dmem.rdata,
    (wb_sel === WB_PC)  -> pc_plus4,
    (wb_sel === WB_CSR) -> csr_rdata
  ))

  when(rf_wen === REN_S) {
    regfile(wb_addr) := wb_data
    printf(p"LW/ADD WriteBack: wb_addr = ${wb_addr}, wb_data = ${wb_data}, regfile(wb_addr) = ${regfile(wb_addr)} !!!\n")
  }

  //**********************************
  // Debug
  io.gp   := regfile(3)
  io.exit := (pc_reg === 0x30.U(WORD_LEN.W))
  printf(p"io.pc      : 0x${Hexadecimal(pc_reg)}\n")
  printf(p"inst       : 0x${Hexadecimal(inst)},  inst_code: ${inst_code} (01 LW, 02 SW, 03 ADD)! \n")
  printf(p"gp         : ${regfile(3)}\n")
  printf(p"rs1_addr(ADD.src1)(LW.base)(SW.base)     : $rs1_addr\n")
  printf(p"rs2_addr(ADD.src2)(LW.offset)(SW.src)    : $rs2_addr\n")
  printf(p"wb_addr(ADD.dest)(LW.dest)(SW.offset)    : $wb_addr\n")
  printf(p"rs1_data   : 0x${Hexadecimal(rs1_data)}\n")
  printf(p"rs2_data   : 0x${Hexadecimal(rs2_data)}\n")
  printf(p"wb_data    : 0x${Hexadecimal(wb_data)}\n")
  printf(p"dmem.addr  : ${io.dmem.addr}\n")
  printf(p"dmem.rdata : ${io.dmem.rdata}\n")
  printf(p"**** regfiles: ${regfile(4)} ${regfile(6)} ${regfile(8)} ${regfile(10)} ${regfile(12)} ${regfile(14)} ****\n")
  printf("---------END\n\n")
}

//------------Top-------------
class Top extends Module {
  val io = IO(new Bundle {
    val exit = Output(Bool())
    val gp   = Output(UInt(WORD_LEN.W))
  })
  val core = Module(new Core())
  val memory = Module(new Memory())
  core.io.imem <> memory.io.imem
  core.io.dmem <> memory.io.dmem
  io.exit := core.io.exit
  io.gp   := core.io.gp
}

//------------GenSvTest-------------
object GenSvTest extends App {
    println("------Begin of emitSystemVerilog......")
    val verilog = _root_.circt.stage.ChiselStage.emitSystemVerilog(new Top);
    println(verilog);
    println("------Enddd of emitSystemVerilog......")
}

